import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4.v1 import p4runtime_pb2
import grpc
import argparse
import json
import os
import threading

# 定义 P4Info 文件和 gRPC 地址
def program_switch(addr, device_id, sw_conf_file, workdir, proto_dump_fpath, runtime_json):
    sw_conf = json.load(sw_conf_file)
    try:
        p4info_fpath = os.path.join(workdir, sw_conf['p4info'])
        p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_fpath)

        target = sw_conf['target']

        print("Connecting to P4Runtime server on %s (%s)..." % (addr, target))

        if target == "bmv2":
            sw = p4runtime_lib.bmv2.Bmv2SwitchConnection(address=addr, device_id=device_id,
                                                         proto_dump_file=proto_dump_fpath)
        else:
            raise Exception("Don't know how to connect to target %s" % target)

        try:
            sw.MasterArbitrationUpdate()

            if target == "bmv2":
                print("Setting pipeline config (%s)..." % sw_conf['bmv2_json'])
                bmv2_json_fpath = os.path.join(workdir, sw_conf['bmv2_json'])
                sw.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                               bmv2_json_file_path=bmv2_json_fpath)
            else:
                raise Exception("Should not be here")

            if 'table_entries' in sw_conf:
                table_entries = sw_conf['table_entries']
                print("Inserting %d table entries..." % len(table_entries))
                for entry in table_entries:
                    print(p4runtime_lib.helper.tableEntryToString(entry))
                    insertTableEntry(sw, entry, p4info_helper)

            if 'multicast_group_entries' in sw_conf:
                group_entries = sw_conf['multicast_group_entries']
                print("Inserting %d group entries..." % len(group_entries))
                for entry in group_entries:
                    print(p4runtime_lib.helper.groupEntryToString(entry))
                    insertMulticastGroupEntry(sw, entry, p4info_helper)

            if 'clone_session_entries' in sw_conf:
                clone_entries = sw_conf['clone_session_entries']
                print("Inserting %d clone entries..." % len(clone_entries))
                for entry in clone_entries:
                    print(p4runtime_lib.helper.cloneEntryToString(entry))
                    insertCloneGroupEntry(sw, entry, p4info_helper)

        finally:
            sw.shutdown()
    except Exception as e:
        print("Error: %s" % str(e))


def insertTableEntry(sw, entry, p4info_helper):
    table_entry = p4info_helper.buildTableEntry(
        table_name=entry['table'],
        match_fields=entry.get('match', {}),
        action_name=entry['action_name'],
        action_params=entry.get('action_params', {})
    )
    sw.WriteTableEntry(table_entry)


def insertMulticastGroupEntry(sw, entry, p4info_helper):
    multicast_group_entry = p4info_helper.buildMulticastGroupEntry(
        multicast_group_id=entry['multicast_group_id'],
        replicas=entry['replicas']
    )
    sw.WriteMulticastGroupEntry(multicast_group_entry)


def insertCloneGroupEntry(sw, entry, p4info_helper):
    clone_session_entry = p4info_helper.buildCloneSessionEntry(
        clone_session_id=entry['clone_session_id'],
        packet_length=entry.get('packet_length', 128),
        replicas=entry['replicas']
    )
    sw.WriteCloneSessionEntry(clone_session_entry)


# 定义接收克隆数据包的回调函数
def receive_clone_packets(stub, device_id):
    request = p4runtime_pb2.StreamMessageRequest()
    request.arbitration.device_id = device_id
    request.arbitration.election_id.high = 0
    request.arbitration.election_id.low = 1

    # 发送仲裁请求
    stub.StreamChannel(request)

    # 接收克隆数据包
    for response in stub.StreamChannel(request):
        if response.HasField('packet'):
            packet = response.packet
            # 这里需要根据实际的数据包格式解析 ECN 值
            # 假设 ECN 值在 IPv4 头部的特定位置
            ecn_value = packet.payload[13] & 0x03  # 假设 ECN 在第 13 字节的低 2 位
            if ecn_value == 3:
                print("Network congestion detected!")
                # 可以在这里下发新的运行时规则
                # 示例：修改队列阈值
                new_threshold = 50
                table_entry = p4info_helper.buildTableEntry(
                    table_name='queue_threshold_table',
                    match_fields={
                        '': ()
                    },
                    action_name='set_queue_threshold',
                    action_params={
                        'threshold': new_threshold
                    }
                )
                request = p4runtime_pb2.WriteRequest()
                request.device_id = device_id
                update = request.updates.add()
                update.type = p4runtime_pb2.Update.MODIFY
                update.entity.table_entry.CopyFrom(table_entry)
                stub.Write(request)
                print(f"New queue threshold {new_threshold} has been set.")


def main():
    parser = argparse.ArgumentParser(description='P4Runtime Simple Controller')

    parser.add_argument('-a', '--p4runtime-server-addr',
                        help='address and port of the switch\'s P4Runtime server (e.g. 192.168.0.1:50051)',
                        type=str, action="store", required=True)
    parser.add_argument('-d', '--device-id',
                        help='Internal device ID to use in P4Runtime messages',
                        type=int, action="store", required=True)
    parser.add_argument('-p', '--proto-dump-file',
                        help='path to file where to dump protobuf messages sent to the switch',
                        type=str, action="store", required=True)
    parser.add_argument("-c", '--runtime-conf-file',
                        help="path to input runtime configuration file (JSON)",
                        type=str, action="store", required=True)

    args = parser.parse_args()

    if not os.path.exists(args.runtime_conf_file):
        parser.error("File %s does not exist!" % args.runtime_conf_file)
    workdir = os.path.dirname(os.path.abspath(args.runtime_conf_file))
    with open(args.runtime_conf_file, 'r') as sw_conf_file:
        program_switch(addr=args.p4runtime_server_addr,
                       device_id=args.device_id,
                       sw_conf_file=sw_conf_file,
                       workdir=workdir,
                       proto_dump_fpath=args.proto_dump_file,
                       runtime_json=args.runtime_conf_file)

    # 建立 gRPC 连接
    channel = grpc.insecure_channel(args.p4runtime_server_addr)
    stub = p4runtime_pb2.P4RuntimeStub(channel)

    # 启动接收克隆数据包的线程
    clone_thread = threading.Thread(target=receive_clone_packets, args=(stub, args.device_id))
    clone_thread.start()

    print("Waiting for clone packets...")


if __name__ == '__main__':
    main()